import matplotlib.pyplot as plt
import numpy as np


# 从文件中读取损失数据
def read_loss_data(file_path):
    losses = []
    with open(file_path, "r") as file:
        for line in file:
            if "loss in epoch" in line:
                # 提取损失值
                loss = float(line.strip().split(": ")[-1])
                losses.append(loss)
    return losses


# 读取数据
# TODO 这里需要改地址
routing_transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/sasrec_ex/myCode/ml-1m_output/routingTransformer_20250102_144322.txt"
)
transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/sasrec_ex/myCode/ml-1m_output/transformer-output.txt"
)

# 绘制损失对比图
plt.figure(figsize=(10, 5))
plt.plot(routing_transformer_losses, label="Routing Transformer", alpha=0.7)
plt.plot(transformer_losses, label="Transformer", alpha=0.7)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.title("Loss Comparison between Routing Transformer and Transformer")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig("ml-1m_output/loss_comparison.png")

# 绘制运行时间和准确性的对比图
plt.figure(figsize=(18, 12))

# 绘制损失对比图
plt.subplot(3, 1, 1)
plt.plot(routing_transformer_losses, label="Routing Transformer", alpha=0.7)
plt.plot(transformer_losses, label="Transformer", alpha=0.7)
plt.xlabel("Iteration")
plt.ylabel("Loss")
# TODO 这里改时间
plt.title("Loss Comparison 1-2")
plt.legend()
plt.grid(True)

# TODO 这里需要改参数，绘制运行时间对比
# epoch:100, time: 1361.599287(s), valid (NDCG@10: 0.5878, HR@10: 0.8242), test (NDCG@10: 0.5633, HR@10: 0.8018)
# epoch:100, time: 1407.427627(s), valid (NDCG@10: 0.5774, HR@10: 0.8131), test (NDCG@10: 0.5536, HR@10: 0.7871)
valid_time_values = [1407.427627, 1796.537418]  # 前面的是Routing
plt.subplot(3, 1, 2)
plt.bar(
    ["Routing Transformer", "Transformer"],
    valid_time_values,
    color=["lightblue", "salmon"],
)
plt.ylabel("Time (seconds)")
plt.title("Running Time Comparison")

# 绘制准确性对比 前面的是Routing
valid_ndcg_values = [0.5774, 0.6117]
valid_hr_values = [0.8131, 0.8455]
test_ndcg_values = [0.5536, 0.5838]
test_hr_values = [0.7871, 0.8157]
plt.subplot(3, 1, 3)
bar_width = 0.2
index = np.arange(len(valid_ndcg_values))

plt.bar(
    index,
    valid_ndcg_values,
    bar_width,
    label="Valid NDCG@10",
    color="lightgreen",
    alpha=0.7,
)
plt.bar(
    index + bar_width,
    valid_hr_values,
    bar_width,
    label="Valid HR@10",
    color="orange",
    alpha=0.7,
)
plt.bar(
    index + bar_width * 2,
    test_ndcg_values,
    bar_width,
    label="Test NDCG@10",
    color="darkgreen",
    alpha=0.7,
)
plt.bar(
    index + bar_width * 3,
    test_hr_values,
    bar_width,
    label="Test HR@10",
    color="red",
    alpha=0.7,
)
plt.xlabel("Model")
plt.ylabel("Score")
plt.title("Accuracy Metrics Comparison")
plt.xticks(index + bar_width * 1.5, ["Routing Transformer", "Transformer"])
plt.legend()

plt.tight_layout()
plt.savefig("ml-1m_output/enhanced_performance_comparison_1-2.png")
plt.show()
